package com.eonmind.ai.core.chat.platform;

import com.eonmind.ai.service.impl.AiPlatformService;
import com.eonmind.common.contract.enums.ErrorCode;
import com.eonmind.common.contract.enums.StatusEnum;
import com.eonmind.common.contract.exception.BusinessException;
import com.eonmind.common.provider.entity.ai.AiPlatform;
import com.eonmind.common.provider.entity.ai.AiPlatformModel;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.CommandLineRunner;
import org.springframework.stereotype.Component;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

@Slf4j
@Component
public class AiPlatformCache implements CommandLineRunner {

    @Resource
    private AiPlatformService aiPlatformService;

    private static final ConcurrentMap<Long, List<AiPlatformModel>> modelCache = new ConcurrentHashMap<>();
    private static final ConcurrentMap<Long, AtomicInteger> modelCountCache = new ConcurrentHashMap<>();
    private final Runnable task = () -> {
        log.info("AI Platform Loading...");
        // 获取所有模型与平台的绑定
        List<AiPlatformModel> aiPlatformModelList = aiPlatformService.getPlatformModelByStatus(StatusEnum.OPENED.getCode());

        // 获取所有平台的ID
        Set<Long> platformIds = aiPlatformModelList.stream()
                .map(AiPlatformModel::getPlatformId)
                .collect(Collectors.toSet());
        List<AiPlatform> aiPlatformList = aiPlatformService.getByIds(platformIds);
        Map<Long, AiPlatform> aiPlatformMap = aiPlatformList.stream()
                .collect(Collectors.toMap(AiPlatform::getId, v -> v));
        log.info("Loaded {} AI Platforms", aiPlatformModelList.size());

        // 按照模型id分组
        Map<Long, List<AiPlatformModel>> modelMap = aiPlatformModelList.stream()
                .collect(Collectors.groupingBy(AiPlatformModel::getModelId));
        modelMap.forEach((modelId, modelList) -> {
            modelCountCache.put(modelId, new AtomicInteger(0));
            modelList.sort(Comparator.comparing(AiPlatformModel::getPriority));
        });
        modelCache.putAll(modelMap);

        modelMap.forEach((modelId, modelList) -> {
            modelCountCache.put(modelId, new AtomicInteger());
        });
        log.info("Load AI Platforms Success");
    };


    @Override
    public void run(String... args) throws Exception {
        log.info("AI Platform Loading...");
        Executors.newScheduledThreadPool(1).scheduleAtFixedRate(task, 0, 1, TimeUnit.MINUTES);
    }

    public AiPlatform getPlatform(Long modelId) {
        AtomicInteger atomicInteger = modelCountCache.get(modelId);
        if (atomicInteger == null) {
            throw new BusinessException(ErrorCode.FAIL, "模型不存在可用的渠道");
        }
        int andAdd = atomicInteger.getAndAdd(1);
        int size = modelCache.get(modelId).size();
        AiPlatformModel aiPlatformModel = modelCache.get(modelId).get(andAdd % size);
        if (aiPlatformModel == null) {
            throw new BusinessException(ErrorCode.FAIL, "模型不存在可用的渠道");
        }
        Long platformId = aiPlatformModel.getPlatformId();
        AiPlatform platformById = aiPlatformService.getPlatformById(platformId);
        if (platformById == null) {
            throw new BusinessException(ErrorCode.FAIL, "模型不存在可用的渠道");
        }
        return platformById;
    }


}
